import json
import os
from traceback import print_tb
from jax import grad
import jax.numpy as jnp
from jax import jit
import time
from matplotlib import pyplot as plt
import numpy as np
import numpy.random as npr
import jax
import jax.numpy as jnp
from jax import device_put
from jax import jit, grad, lax, random
from jax.example_libraries import optimizers
from jax.example_libraries import stax
from jax.example_libraries.stax import Dense, FanOut, Relu, Softplus, Sigmoid, FanInSum
from jax.nn import sigmoid
from functools import partial
from jax import vmap
from flax import linen as nn
from flax.training import train_state
from flax import struct
from jax import lax

from jax import tree_util
from jax.tree_util import tree_structure
from jax.tree_util import tree_flatten, tree_unflatten

import jax.experimental.sparse as sparse

import optax
import cv2

from openTSNE import TSNE

def f(x):
   return jnp.asarray(
     [x[0], 5*x[2], 4*x[1]**2 - 2*x[2], x[2] * jnp.sin(x[0])])

# 计算 jacobian matrix
print(jax.jacrev(f)(jnp.array([1., 2., 3.])))

exit()

rnn_log_pth = "./data/adaptive_trajectory_optimization/rnn_state_opt_log.npy"
trj_log_pth = "./data/adaptive_trajectory_optimization/trj_log.json"

rnn_log_pth = "./data/adaptive_trajectory_optimization/rnn_state_opt_log_32.npy"
trj_log_pth = "./data/adaptive_trajectory_optimization/trj_log_32.json"

# load ./data/adaptive_trajectory_optimization/rnn_state_opt_log.npy
rnn_state_opt_log = np.load(rnn_log_pth)

print("shape of rnn_state_opt_log: ", rnn_state_opt_log.shape)

# # # plot each element of rnn_state_opt_log
# # for i in range(len(rnn_state_opt_log)):
# #     plt.bar(np.arange(0, 128), rnn_state_opt_log[i])
# #     plt.show()
# # id1 = 19
# # id2 = 29
# # id1 = 21
# # id2 = 27
# # id1 = 25
# # id2 = 30
# # id1 = 22
# # id2 = 60
# id1 = 22
# id2 = 23
# norm_id1 = rnn_state_opt_log[id1] / np.linalg.norm(rnn_state_opt_log[id1])
# norm_id2 = rnn_state_opt_log[id2] / np.linalg.norm(rnn_state_opt_log[id2])
# print("similarity : ", np.dot(norm_id1, norm_id2))
# plt.bar(np.arange(0, 128), rnn_state_opt_log[id1], color='r')
# plt.bar(np.arange(0, 128), rnn_state_opt_log[id2], color='g')
# plt.show()

# erase the first element of rnn_state_opt_log
rnn_state_opt_log = rnn_state_opt_log[1:]

# # perform t-SNE on rnn_state_opt_log
# tsne_embedding = TSNE().fit(rnn_state_opt_log)
# print("shape of tsne_embedding: ", tsne_embedding.shape)
# # visualize the t-SNE embedding
# plt.scatter(tsne_embedding[:, 0], tsne_embedding[:, 1])
# plt.show()

# compute the self similarity matrix of rnn_state_opt_log
sim_mat = np.zeros((len(rnn_state_opt_log), len(rnn_state_opt_log)))
for i in range(len(rnn_state_opt_log)):
    for j in range(len(rnn_state_opt_log)):
        norm_i = rnn_state_opt_log[i] / np.linalg.norm(rnn_state_opt_log[i])
        norm_j = rnn_state_opt_log[j] / np.linalg.norm(rnn_state_opt_log[j])
        sim_mat[i, j] = np.dot(norm_i, norm_j)

# create a matrix with each element being (i, j, sim_value) where i and j are the indices of the element and sim_value is the similarity value
sim_mat_3d = np.zeros((len(rnn_state_opt_log), len(rnn_state_opt_log), 3))
for i in range(len(rnn_state_opt_log)):
    for j in range(len(rnn_state_opt_log)):
        sim_mat_3d[i, j, 0] = i
        sim_mat_3d[i, j, 1] = j
        sim_mat_3d[i, j, 2] = sim_mat[i, j]

# sort the matrix by the similarity value
sim_mat_3d = sim_mat_3d.reshape(-1, 3)
sim_mat_3d = sim_mat_3d[sim_mat_3d[:, 2].argsort()]
sim_mat_3d = sim_mat_3d[::-1]

# select the top 100 elements except the elements with the same i and j
top_100 = []
for i in range(len(sim_mat_3d)):
    if len(top_100) >= 200:
        break
    if sim_mat_3d[i, 0] != sim_mat_3d[i, 1]:
        top_100.append(sim_mat_3d[i])

# print("top 100 similarity: ", top_100)

with open(trj_log_pth, "r") as f:
    trj_json = json.load(f)
    trj_log = trj_json["data"]

# draw each element of trj_log one single opencv canvas
img_seq = []
for i in range(len(trj_log)):
    img = np.zeros((60, 60, 3), np.uint8)
    scale = 4
    for j in range(len(trj_log[i])-1):
        cv2.line(img, (int(trj_log[i][j][0]/scale), int(trj_log[i][j][1]/scale)), (int(trj_log[i][j+1][0]/scale), int(trj_log[i][j+1][1]/scale)), (255, 255, 255), 1)
    # draw a bigger dot at the start point
    cv2.circle(img, (int(trj_log[i][0][0]/scale), int(trj_log[i][0][1]/scale)), 2, (255, 255, 255), -1)
    # draw a white border
    cv2.rectangle(img, (0, 0), (59, 59), (255, 255, 255), 1)
    img_seq.append(img)

n_img = len(img_seq)

# rearrange img_seq on a big canvas with a grid formation
n_row = 10
n_col = 10
big_img = np.zeros((n_row*60, n_col*60, 3), np.uint8)
img_id = 0
for i in range(n_row):
    for j in range(n_col):
        if img_id < n_img:
            big_img[i*60:(i+1)*60, j*60:(j+1)*60, :] = img_seq[img_id]
        img_id += 1

# detect the mouse position and print which grid it is in
def mouse_callback(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        print("x: ", x, "y: ", y)
        print("grid: ", int(x/60), int(y/60))
        print("img_id: ", int(y/60)*n_col + int(x/60))

cv2.namedWindow("big_img")
cv2.setMouseCallback("big_img", mouse_callback)
cv2.imshow("big_img", big_img)
cv2.waitKey(0)

for i in range(len(img_seq)):
    big_img_cpy = big_img.copy()
    # draw a green border around the selected image
    cv2.rectangle(big_img_cpy, (int(i%n_col)*60, int(i/n_col)*60), (int(i%n_col)*60+59, int(i/n_col)*60+59), (0, 255, 0), 2)
    siblings = []
    for j in range(len(top_100)):
        if (top_100[j][0] == i):
            siblings.append(top_100[j][1])
    for s in siblings:
        # draw a red border around the sibling
        cv2.rectangle(big_img_cpy, (int(s%n_col)*60, int(s/n_col)*60), (int(s%n_col)*60+59, int(s/n_col)*60+59), (0, 0, 255), 2)
    cv2.imshow("big_img", big_img_cpy)
    k = cv2.waitKey(0)
    if k == ord('q'):
        break


# for i in range(len(top_100)):
#     big_img_cpy = big_img.copy()
#     id1 = int(top_100[i][0])
#     id2 = int(top_100[i][1])
#     sim = top_100[i][2]
#     # draw a red border around the first image
#     cv2.rectangle(big_img_cpy, (int(id1%n_col)*60, int(id1/n_col)*60), (int(id1%n_col)*60+59, int(id1/n_col)*60+59), (0, 0, 255), 2)
#     # draw a green border around the second image
#     cv2.rectangle(big_img_cpy, (int(id2%n_col)*60, int(id2/n_col)*60), (int(id2%n_col)*60+59, int(id2/n_col)*60+59), (0, 255, 0), 2)
#     # draw similarity value on the bottom left corner of big_img_cpy
#     cv2.putText(big_img_cpy, "sim: " + str(sim), (10, 590), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
#     cv2.imshow("big_img", big_img_cpy)
#     k = cv2.waitKey(0)
#     if k == ord('q'):
#         break

